import os
import xml.etree.ElementTree as ET
from typing import List, Tuple

from .base import Base

VOC_CLASS_LIST = [
    "aeroplane",
    "bicycle",
    "boat",
    "bus",
    "car",
    "motorbike",
    "train",
    "bottle",
    "chair",
    "diningtable",
    "pottedplant",
    "sofa",
    "tvmonitor",
    "bird",
    "cat",
    "cow",
    "dog",
    "horse",
    "sheep",
    "person",
]


class VOCUtils(Base):
    def __init__(
        self, root_path: str, version: List[str], auto_download: bool = True
    ) -> None:
        self.class_list = VOC_CLASS_LIST
        super().__init__(root_path, version, auto_download)

    def _prepare_data(self, version: str, auto_download: bool) -> None:
        if not os.path.exists(self.root_path):
            os.mkdir(self.root_path)
        for v in version:
            assert v in ["2007", "2012"], f"Unsupported version {v}"
            save_path = os.path.join(self.root_path, f"VOC{v}")

            if not os.path.exists(save_path):
                if os.path.exists(save_path + ".tar"):
                    pass
                elif auto_download:
                    self._download(v)
                else:
                    raise FileNotFoundError

                import shutil
                import tarfile

                with tarfile.open(save_path + ".tar", "r") as tar:
                    tar.extractall(path=self.root_path)

                shutil.move(
                    os.path.join(self.root_path, "VOCdevkit", f"VOC{v}"),
                    self.root_path,
                )
                os.remove(save_path + ".tar")
                shutil.rmtree(os.path.join(self.root_path, "VOCdevkit"))
                os.mkdir(os.path.join(save_path, "Yolo"))

    def _download(self, version: str) -> None:

        import requests
        from tqdm import tqdm

        url = {
            "2007": "http://host.robots.ox.ac.uk/pascal/VOC/voc2007/VOCtrainval_06-Nov-2007.tar",
            "2012": "http://host.robots.ox.ac.uk/pascal/VOC/voc2012/VOCtrainval_11-May-2012.tar",
        }

        save_path = os.path.join(self.root_path, f"VOC{version}")

        if not os.path.exists(save_path + ".tar"):
            try:
                response = requests.get(
                    url[version],
                    stream=True,
                )
                response.raise_for_status()
                total_size = int(response.headers.get('content-length', 0))

                with open(save_path + ".tar", "ab") as f, tqdm(
                    desc=f"Downloading dataset to {save_path}",
                    total=total_size,
                    unit='B',
                    unit_scale=True,
                    unit_divisor=1024,
                ) as bar:
                    for chunk in response.iter_content(chunk_size=1024):
                        size = f.write(chunk)
                        bar.update(size)
            except requests.exceptions.ChunkedEncodingError:
                os.remove(save_path + ".tar")
                raise ConnectionError(
                    f"Download failed, please download from {url[version]} to {save_path}.tar manually and rerun the program."
                )

    def _convert(self, xml_path: str, save_path: str) -> None:
        tree = ET.parse(xml_path)
        root = tree.getroot()
        size = root.find('size')
        image_width = int(size.find('width').text)
        image_height = int(size.find('height').text)
        with open(save_path, "w") as f:
            for obj in root.iter("object"):
                class_name = obj.find("name").text
                class_id = self.class_list.index(class_name)
                bndbox = obj.find("bndbox")
                x_min = int(bndbox.find("xmin").text)
                y_min = int(bndbox.find("ymin").text)
                x_max = int(bndbox.find("xmax").text)
                y_max = int(bndbox.find("ymax").text)

                x_center = (x_min + x_max) / 2
                y_center = (y_min + y_max) / 2
                width = x_max - x_min
                height = y_max - y_min

                x_center /= image_width
                y_center /= image_height
                width /= image_width
                height /= image_height

                f.write(f"{class_id} {x_center} {y_center} {width} {height}\n")

    def _get_image_label_list(
        self, version: List[str], is_train: bool
    ) -> Tuple[List[str], List[str], List[str], List[str]]:
        images_path, labels_path = (
            [],
            [],
        )
        mode = "train" if is_train else "val"
        for version in version:
            root_path = os.path.join(self.root_path, f"VOC{version}")
            with open(
                os.path.join(root_path, "ImageSets", "Main", f"{mode}.txt"), "r"
            ) as f:
                idx_list = [line.strip() for line in f.readlines()]

            for filename in idx_list:
                images_path.append(
                    os.path.join(root_path, "JPEGImages", f"{filename}.jpg")
                )

                if not os.path.exists(
                    os.path.join(root_path, "Yolo", f"{filename}.txt")
                ):
                    self._convert(
                        os.path.join(root_path, "Annotations", f"{filename}.xml"),
                        os.path.join(root_path, "Yolo", f"{filename}.txt"),
                    )
                labels_path.append(os.path.join(root_path, "Yolo", f"{filename}.txt"))

        return images_path, labels_path
